package OrderedCoverPackage;

import EstimatorPackage.Expressions.ComplexExpression;
import EstimatorPackage.Expressions.EmptyExpression;
import EstimatorPackage.Expressions.Expression;
import EstimatorPackage.Expressions.ListExpression;
import EstimatorPackage.Expressions.SimpleExpression;
import DataSetPackage.View;
import EstimatorPackage.AbstractEstimatorExp;
import EstimatorPackage.Consts;
import EstimatorPackage.EstimatorExp;
import HierarchyPackage.HClass;
import HierarchyPackage.Hierarchy;
import HierarchyPackage.ProductHierarchy.ProductHClass;
import HierarchyPackage.ProductHierarchy.ProductHierarchy;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 *
 * @author Martin Haulrich
 */
public class OrderedCoverEstimator extends AbstractOC implements EstimatorExp {

    final Expression empty;

//    final ClassAlgebraSingleton CA;
    Map<OCPartition, Expression> ocExpr;

    Map<HClass, Double> classCache;

    public OrderedCoverEstimator(BaseHPM cover) {
        super(cover.getHierarchy(), cover.prior);

        partitionList = new LinkedList<OCPartition>();
//        CA = ClassAlgebraSingleton.getInstance();
        OCPartition ctop = cover.getTop();
        top = new OCPartition(ctop.mainClass, ctop.ID);
        top.restLogEmpProb = ctop.restLogEmpProb;
        top.restLogPriorProbDenom = ctop.restLogPriorProbDenom;
        top.resLogPseuEmpProb = ctop.resLogPseuEmpProb;

        partitionList.add(top);
        ocExpr = new HashMap<OCPartition, Expression>(cover.partitionList().size());

//        System.err.println(top);
//        System.err.println(partitionList);
//        System.err.println(partitionList.indexOf(top));
//        System.err.println(partitionList.contains(top) + "\t"+ partitionList.contains(ctop));
//        System.err.println(cover.getTop().children);
//        System.err.println(ctop.children);
        buildHierarchyRec(ctop, top);

        empty = new EmptyExpression();
        //   makeMap();
        classCache = new LinkedHashMap<HClass, Double>(25600, 0.75F, true) {

            @Override
            protected boolean removeEldestEntry(Map.Entry<HClass, Double> eldest) {
                return size() > 20000;
            }
        };

    }

    public OrderedCoverEstimator(String fn, ProductHierarchy hier, EstimatorExp prior) throws IOException {
        super(hier, prior);
        empty = new EmptyExpression();
//        CA = ClassAlgebraSingleton.getInstance();
        ocExpr = new HashMap<OCPartition, Expression>();
        load(fn);
        classCache = new LinkedHashMap<HClass, Double>(12800, 0.75F, true) {

            @Override
            protected boolean removeEldestEntry(Map.Entry<HClass, Double> eldest) {
                return size() > 10000;
            }
        };

    }

    private void buildHierarchyRec(OCPartition trainPart, OCPartition myPart) {

        //      System.err.println("Adding: " + trainPart.mainClass.getName());
        //    System.err.println(trainPart.children);
        if (trainPart.children == null) {
            //        System.err.println("NO CHILDREN");
            myPart.children = null;
            return;
        }
        List<OCPartition> myChildren = new ArrayList<OCPartition>(trainPart.children.size());
        for (OCPartition trCh : trainPart.children) {
            OCPartition myCh = new OCPartition(trCh.mainClass, trCh.ID);
            myCh.parent = myPart;
            myCh.restLogEmpProb = trCh.restLogEmpProb;
            myCh.resLogPseuEmpProb = trCh.resLogPseuEmpProb;
            myCh.restLogPriorProbDenom = trCh.restLogPriorProbDenom;
            myChildren.add(myCh);
            myCh.ghost = trCh.ghost;
        }
        //System.err.println(myChildren);
        myPart.children = myChildren;
//        System.err.println(myPart.mainClass.getName());
//        System.err.println(partitionList);
//        for (OCPartition oc : partitionList) {
//            System.err.println(oc+ "\t" + oc.equals(myPart));
//            System.err.println(oc+ "\t" + oc.equals(trainPart));
//        }
//        System.err.println(myPart);
        partitionList.addAll(partitionList.indexOf(myPart), myChildren);

        for (int i = 0; i < myChildren.size(); i++) {
            OCPartition trCh = trainPart.children.get(i);
            OCPartition myCh = myChildren.get(i);
            buildHierarchyRec(trCh, myCh);
        }


    }

    @Override
    public double getEstimate(Integer did, View v) {
//        ProductHClass didClass = (ProductHClass) hierarchy.newClass(did, v);
//        System.err.println("ClassPriorProb: " + getEstimate(didClass));


        OCPartition p = findPartition(did, v, top);

        System.err.println("        " + p.mainClass.getName());

        double priorProb = prior.getEstimate(did, v);

        //     System.err.println("priorProb: " + priorProb + "\tpriorProbN: " + priorProbN);

        //System.err.print(p.mainClass.getName() + "\t");
        double w = getLogWeight(p);
        //         System.err.println("w: " + w);
        //System.err.println((priorProb + w) + " = " + priorProb + " + " + w);
        double prob = priorProb + w;

        return prob;
    }

    public double classAsDataPointEstimate(HClass c) {

        OCPartition p = findFirstPartition(c, top);

        double priorProb = prior.getEstimate(new SimpleExpression(c), hierarchy);
        double w = getLogWeight(p);
        double prob = priorProb + w;
//        System.err.println(prob + " = " + priorProb + " + " + w);
        return prob;
    }

    public double getEstimate(HClass c) {

//        System.err.println("------ " + c.getName() + " -------------");
//        System.err.println("PL-size: " + partitionList.size());

        Double est = classCache.get(c);
        if (est == null) {
            Expression cexp = new SimpleExpression(c);

            double sum = Consts.logZero;
            List<OCPartition> parts = new ArrayList<OCPartition>();
            findPartitionsRec(c, top, parts);
//            System.err.println("parts-size: " + parts.size());
//            System.err.println("Partition-list-size: " + partitionList.size());
            for (OCPartition p : parts) {
                if (!p.ghost) {
                    double w = getLogWeight(p);
                    Expression pmexp = p.mainExp;
                    Expression psexp = ocExpr.get(p);
                    if (psexp == null) {
////                        psexp = CA.map(new ListExpression(ListExpression.UNION, subtractsExpression(p)), hierarchy);
                        psexp = AbstractEstimatorExp.reduceList(new ListExpression(ListExpression.UNION, subtractsExpression(p)), hierarchy);
                        ocExpr.put(p, psexp);
                    }
//                    Expression psexp = new ListExpression(ListExpression.UNION, subtractsExpression(p));

                    Expression pexp = new ComplexExpression(ComplexExpression.DIFFERENCE, pmexp, psexp);
                    Expression intersect = new ComplexExpression(ComplexExpression.INTERSECTION, pexp, cexp);

                    double priorProb = prior.getEstimate(intersect, hierarchy);
//                    System.err.println("int: " + intersect);
//System.err.println(p.mainExp + " : " + priorProb + "\t" + prior.getEstimate(pexp, hierarchy));

                    if (Double.isNaN(priorProb)) {
                        priorProb = Consts.logZero;
                    } else {
                        if (priorProb > Consts.logZero) {

//                            System.err.println("s1: " + sum);
                            double pr = priorProb + w;
//                            System.err.println(Math.pow(10, pr) + " --- " + Math.pow(10, sum));
//                            System.err.println(p.mainExp + "\t" + priorProb + " + " + w + " = " + pr);
                            double l = Math.max(sum, pr);
//                            System.err.println("l: " + l + " + " + (Math.pow(10, sum - l)) + " + "+(Math.pow(10, pr - l)));
//                            System.err.println("l1: " + (l + Math.log10(Math.pow(10, sum - l))));
//                            System.err.println("l2: " + (Math.log10(Math.pow(10, sum - l) + Math.pow(10, pr - l))));
//                            System.err.println("l3: " +  (l + Math.log10(Math.pow(10, sum - l) + Math.pow(10, pr - l))));
                            sum = l + Math.log10(Math.pow(10, sum - l) + Math.pow(10, pr - l));
//                            System.err.println("s2: " + sum);

                        }
                    }
                }
//                System.err.println("sum: " + sum);
            }

            est = sum;
//            System.err.println("est: " + est);
            classCache.put(c, est);
        }
        return est;
    }

    public double getConditionalEstimate(Integer did, View v) {

        double unCond = getEstimate(did, v);

        ProductHClass didClass = (ProductHClass) hierarchy.newClass(did, v);

        return unCond - getConditioningProb(didClass);
    }

    public double getConditioningProb(ProductHClass pc) {

        HClass[] classes = new HClass[pc.getID().length];
        System.arraycopy(pc.getID(), 0, classes, 0, pc.getID().length);
        classes[classes.length - 1] = ((ProductHierarchy) hierarchy).getHierarchies()[classes.length - 1].getTop();
        ProductHClass newTopClass = new ProductHClass((ProductHierarchy) hierarchy, classes);
//        System.err.println(pc.getName() + "\t" + newTopClass.getName());
        return getEstimate(newTopClass);
    }

    /**
     * Return null if this was the wrong place to go - because current is a ghost-class that actually
     * does not hold the current datapoint.
     * @param did
     * @param v
     * @param current
     * @return
     *
     */
    private OCPartition findPartition(Integer did, View v, OCPartition current) {
        if (!(current.children == null)) {

            for (OCPartition ch : current.children) {
                if (hierarchy.isClassifiedBy(did, v, ch.mainClass)) {
                    OCPartition f = findPartition(did, v, ch);
                    if (f != null) {
                        return f;
                    }
                }
            }
        }
        if (current.ghost) {
            return null;
        }
        return current;
    }

    private OCPartition findFirstPartition(HClass c, OCPartition current) {

        if (current.children != null) {


            for (OCPartition ch : current.children) {
                //  System.err.println("--- " + c.getName() + " trying " + ch.mainClass.getName());
                //    if (!hierarchy.getIntersection(ch.mainClass, c).isEmpty()) {
                if (hierarchy.superclassOf(ch.mainClass, c) || ch.mainClass.equals(c)) {
                    //     System.err.println("---------- YES");
                    OCPartition f = findFirstPartition(c, ch);
                    if (f != null) {
                        return f;
                    }
                }
            //  System.err.println("---------- NO");
            }
        }
        if (current.ghost) {
            return null;
        }
        return current;
    }

    private void findPartitionsRec(HClass c, OCPartition current, List<OCPartition> res) {

        if (getChildren(current).isEmpty()) {
            res.add(current);
            return;
        }
        for (OCPartition ch : getChildren(current)) {

            if (!hierarchy.getIntersection(c, current.mainClass).isEmpty()) {
                findPartitionsRec(c, ch, res);
            }
        }

        // For now we just add the rest-partition hoping that this we be filtered out later.
        // When the ClassAlgebra improves we could check if the intersection is empty first
        res.add(current);
    }

    @Override
    public double getEstimate(Expression pe) {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    @Override
    public double getEstimate(Expression pe, Hierarchy hier) {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    @Override
    protected double getLogWeight(OCPartition p) {
//        System.err.println(p.mainClass.getName());
//        System.err.println(p.logEmpProb + "\t" + p.logPriorProbDenom);
//        return p.restLogEmpProb - p.restLogPriorProbDenom;
        return p.resLogPseuEmpProb - p.restLogPriorProbDenom;
    }

    public void p() {
        System.err.println(partitionList);
    }

    public void save(String fn) throws IOException {

        Set<HClass> classes = new HashSet<HClass>();
        for (OCPartition p : partitionList) {
            classes.add(p.mainClass);
        }

        Map<HClass, Integer> class2id = hierarchy.saveEst(fn + "-hier", classes);

        PrintWriter out = new PrintWriter(new FileWriter(fn));
        for (OCPartition p : partitionList) {
            out.print(p.ID + " ");
        }
        out.println("");

        for (OCPartition p : partitionList) {
            out.println(p.ID + "\t" + class2id.get(p.mainClass) + "\t" + p.mainExp + "\t" + p.restLogEmpProb + "\t" + p.resLogPseuEmpProb + "\t" + p.restLogPriorProbDenom + "\t" + p.ghost);
        }

        out.println("");

        for (OCPartition p : partitionList) {
            if (!(p.children == null || p.children.isEmpty())) {
                out.print(p.ID);
                out.print(" -> ");
                for (OCPartition ch : p.children) {
                    out.print(ch.ID + " ");
                }
                out.println("");
            }

        }

        out.close();
    }

    public void load(String fn) throws IOException {

        BufferedReader input = new BufferedReader(new FileReader(fn));
        Set<Integer> ids = new HashSet<Integer>();

        List<String[]> sl = new ArrayList<String[]>();
        String line = input.readLine();
        String[] pArray = line.split(" ");
//        System.err.println(Arrays.toString(pArray));
        List<String> pList = Arrays.asList(pArray);
//        System.err.println(pList);
        partitionList = new LinkedList<OCPartition>();
        for (String i : pArray) {
            partitionList.add(null);
        }

        while ((line = input.readLine()) != null) {
//            System.err.println(line);
            if (line.equals("")) {
                break;
            } else {
                String tok[] = line.split("\t");
                sl.add(tok);
                ids.add(Integer.parseInt(tok[1]));
            }
        }

        Map<Integer, HClass> id2class = hierarchy.loadEst(fn + "-hier", ids);
//        for (Entry<Long, HClass> e : id2class.entrySet()) {
//            System.err.println(e.getKey() + " : " + e.getValue().getName());
//        }

        Map<Integer, OCPartition> id2partition = new HashMap<Integer, OCPartition>();
        for (String[] t : sl) {
//                 System.err.println(Arrays.toString(t));
            int id = Integer.parseInt(t[0]);
            Integer cid = Integer.parseInt(t[1]);
//            System.err.println("cid: " + cid);
            HClass mainclass = id2class.get(cid);
//            System.err.println("    class: " + mainclass.getName());
            double logEmpProb = Double.parseDouble(t[3]);
            double pseuLogEmpProb = Double.parseDouble(t[4]);
            double logPriorProbDenom = Double.parseDouble(t[5]);
            //            System.err.println("Adding " + id + " with mc: " + mainclass.getName());
            OCPartition p = new OCPartition(mainclass, id);
            p.ghost = Boolean.parseBoolean(t[6]);
            p.restLogEmpProb = logEmpProb;
            p.resLogPseuEmpProb = pseuLogEmpProb;
            p.restLogPriorProbDenom = logPriorProbDenom;
            id2partition.put(id, p);
            partitionList.set(pList.indexOf("" + id), p);
        }

        //       System.err.println(partitionList);
        while ((line = input.readLine()) != null) {
            String[] t = line.split(" -> ");
            int pid = Integer.parseInt(t[0]);
            String[] chs = t[1].split(" ");
            OCPartition parent = id2partition.get(pid);
            List<OCPartition> children = new ArrayList<OCPartition>();
            for (String ch : chs) {
                int chid = Integer.parseInt(ch);
                OCPartition child = id2partition.get(chid);
                child.parent = parent;
                children.add(child);
            }
            parent.children = children;
        }
        top = id2partition.get(0);
        input.close();
    }
}
